import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# print(torch.__version__)

x = torch.rand(512)
w = 3
b = 10
noise = torch.randn(512) * 0.2
y = w * x + b + noise

# print(x)

# 变换特征值的形式
x = torch.unsqueeze(x, dim=1)
y = torch.unsqueeze(y, dim=1)

# print(x)


# 搭建神经网络
class LinearModel(nn.Module):
    def __init__(self, in_fea, out_fea):
        super(LinearModel, self).__init__()
        # 隐藏层
        self.hidden1 = nn.Linear(in_features=in_fea, out_features=6)
        self.hidden2 = nn.Linear(in_features=6, out_features=6)
        self.hidden3 = nn.Linear(in_features=6, out_features=6)
        self.hidden4 = nn.Linear(in_features=6, out_features=6)
        self.hidden5 = nn.Linear(in_features=6, out_features=6)
        self.hidden6 = nn.Linear(in_features=6, out_features=6)

        # 输出层
        self.out = nn.Linear(in_features=6, out_features=out_fea)

    def forward(self, x1):
        x1 = self.hidden1(x1)
        x1 = torch.relu(x1)  # 激活函数进行变换
        x1 = self.hidden2(x1)
        x1 = torch.relu(x1)
        x1 = self.hidden3(x1)
        x1 = torch.relu(x1)
        x1 = self.hidden4(x1)
        x1 = torch.relu(x1)
        x1 = self.hidden5(x1)
        x1 = torch.relu(x1)
        x1 = self.hidden6(x1)
        x1 = torch.relu(x1)
        x1 = self.out(x1)
        return x1


model = LinearModel(1, 1)
loss_func = nn.MSELoss()  # 定义损失函数
optim = torch.optim.SGD(model.parameters(), lr=0.02)  # 定义优化损失的方式(梯度下降)
epoch = 100  # 迭代次数

plt.ion()
for step in range(epoch):
    # 前向传播
    predict = model.forward(x)
    # 计算损失
    loss = loss_func(predict, y)
    # 优化损失
    optim.zero_grad()  # 梯度清零
    loss.backward()  # 反向传播
    optim.step()  # 更新参数
    # print(loss)
    if step % 5 == 0:
        plt.cla()  # 清空画布
        plt.scatter(x.detach().numpy(), y.detach().numpy())  # 真实值-散点图表示
        plt.plot(x.detach().numpy, predict.detach().numpy())  # 预测值-折线图表示
        # 固定画布坐标的范围
        plt.xlim(0, 1.1)
        plt.ylim(0, 20)
        plt.pause(1)

plt.ioff()
plt.show()


















































